package com.fanshuai.websocket.server;

import com.fanshuai.websocket.annotation.WebSocketMapping;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;
import reactor.core.publisher.Mono;

@Service
@WebSocketMapping("/push")
public class WebFluxPushServer implements WebSocketHandler {
    @Override
    public Mono<Void> handle(WebSocketSession session) {
        Flux<String> messages = session.receive()
                .flatMap(webSocketMessage -> {
                    String message = webSocketMessage.getPayloadAsText();
                    System.out.println("from client: " + message);
                    return Mono.just(message);
                });

        return session.send(Flux.create(sink -> {
            for (int i = 0; i < 5; i++) {
                sink.next(session.textMessage("from server: " + i));
            }

            sink.onDispose(messages.map(session::textMessage).subscribe());
        }));
    }
}
